# -*- coding: utf-8 -*-
# @Time    : 2020/11/15 22:00
# @Author  : DuJiabao
# @File    : fshion_mnist.py
# @Description : 
from tensorflow import keras
import tensorflow as tf
import numpy as np

if __name__ == '__main__':
    fashion_mnist = keras.datasets.mnist
    (train_img, train_label), (test_img, test_label) = fashion_mnist.load_data()
    train_img, test_img = train_img / 255., test_img / 255.
    train_img = np.expand_dims(train_img, 3)
    test_img = np.expand_dims(test_img, 3)
    model = keras.Sequential([keras.layers.Conv2D(filters=16, kernel_size=3, strides=1, padding="valid",
                                                  activation=tf.nn.relu, input_shape=[28, 28, 1]),
                              keras.layers.BatchNormalization(),
                              keras.layers.Conv2D(filters=32, kernel_size=3, strides=1, padding="valid",
                                                  activation=tf.nn.relu),
                              keras.layers.BatchNormalization(),
                              keras.layers.Flatten(),
                              keras.layers.Dense(256, activation=tf.nn.relu),
                              keras.layers.Dense(10)]
                             )
    model.compile(optimizer="adam", loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
                  metrics=["accuracy"])
    model.fit(x=train_img, y=train_label, batch_size=128, epochs=3, validation_data=(test_img, test_label))
    model.evaluate(x=test_img, y=test_label, batch_size=128)
    tf.saved_model.save(model, "mnist_model")

    new_model = tf.keras.models.load_model("model")
    # new_model.evaluate(x=test_img, y=test_label, batch_size=128)
